//
// DebugManager.cc
//

#include "DebugManager.h"
#include "DebugFrame.h"

std::map<void *, CDebugManager::AllocListElement *> CDebugManager::_pAllocList;
CDebugManager *CDebugManager::_pDebugManager = NULL;

CDebugManager::CDebugManager(void)
	: _pLog(0)
{
}
	
CDebugManager::~CDebugManager(void)
{
#if defined(_INTERNAL_LEAK_DETECTION)
	if(!_pAllocList.empty()) {
		_writeString(L"[WARNING] Memory leaks detected!\n");
		printMemoryStatus();
		_freeAllocList();
	} else {
		_writeString(L"No memory leaks detected.\n");
	}
#endif // defined(_INTERNAL_LEAK_DETECTION)

	if(_pLog) fclose(_pLog);

	return ;
}

bool CDebugManager::print(const wchar_t *pwszSourcePath,
						 UINT uLineNumber,
						 const wchar_t *pwszFunction,
						 const wchar_t *pwszMessage,
						 ...)
{
	wchar_t buffer[256];
	wchar_t output[256];
	va_list arg;

	va_start(arg, pwszMessage);
	_vsnwprintf(buffer, 256, pwszMessage, arg);
	va_end(arg);

	time_t timeval = time(NULL);
	wchar_t *pTimeStr = _wctime(&timeval);
	pTimeStr[wcslen(pTimeStr) - 1] = '\0';

	const wchar_t *pSource = pwszSourcePath + wcslen(pwszSourcePath) - 1;
	while(pSource != pwszSourcePath) {
		if(*pSource == '\\' || *pSource == '/') break;
		pSource--;
	}

	_snwprintf(output, 256, L"[%s %s:%d] In '%s': %s\n", 
		pTimeStr, pSource, uLineNumber, pwszFunction, buffer);
	_writeString(output);
	_flushStream();

	return true;
}

void CDebugManager::dumpCallStack(void)
{
	const UINT max_name_length = 256;

	CONTEXT context;
	STACKFRAME64 stackframe;
	HANDLE process = 0, thread = 0;
	PSYMBOL_INFO symbol = NULL;
	IMAGEHLP_LINE64 source_info;
	DWORD displacement;
	size_t bufsize;

	// Initialize symbol info structure.
	bufsize = sizeof(SYMBOL_INFO) + (max_name_length - 1) * sizeof(char);
	symbol = static_cast<PSYMBOL_INFO>(::operator new(bufsize));
	if(symbol == NULL) return ;
	memset(symbol, 0, bufsize);
	symbol->SizeOfStruct = sizeof(SYMBOL_INFO);
	symbol->MaxNameLen = max_name_length;

	// Initialize source info structure.
	memset(&source_info, 0, sizeof(IMAGEHLP_LINE64));
	source_info.SizeOfStruct = sizeof(IMAGEHLP_LINE64);
	
	// Setup capture registers.
	RtlCaptureContext(&context);
	memset(&stackframe, 0, sizeof(STACKFRAME64));
	stackframe.AddrPC.Offset = context.Eip;
	stackframe.AddrPC.Mode = AddrModeFlat;
	stackframe.AddrStack.Offset = context.Esp;
	stackframe.AddrStack.Mode = AddrModeFlat;
	stackframe.AddrFrame.Offset = context.Ebp;
	stackframe.AddrFrame.Mode = AddrModeFlat;

	// Get handles.
	process = GetCurrentProcess();
	thread = GetCurrentThread();

	// Initialize dbghelper library.
	if(!SymInitialize(process, NULL, TRUE))
		return ;

	_writeString(L"\n Dumping call stack...\n");
	_writeString(L" Symbol Name      Source File Name                 Line  Address\n");

	// Enumerate each frame inside the stack.
	while(StackWalk64(IMAGE_FILE_MACHINE_I386, process, thread, &stackframe, 
		&context, NULL, SymFunctionTableAccess64, SymGetModuleBase64, NULL))
	{
		// Bottom reached.
		if(stackframe.AddrFrame.Offset == 0)
			break;
		
		// Get symbol name.
		if(SymFromAddr(process, stackframe.AddrPC.Offset, NULL, symbol))
			printRawString(L" %-16.16S ", symbol->Name);

		// Get source code info & function address.
		if(SymGetLineFromAddr64(process, stackframe.AddrPC.Offset, 
			&displacement, &source_info)) {
				int pos = strlen(source_info.FileName);
				while(pos > 0 && source_info.FileName[pos] != '\\') pos--;
				printRawString(L"%-32.32S %-5d 0x%08LX\n", 
					source_info.FileName + pos, 
					source_info.LineNumber,
					stackframe.AddrPC.Offset);
		} else {
			if(GetLastError() == 0x1E7) {
				printRawString(L"No debug symbol loaded.                0x%08LX\n",
					stackframe.AddrPC.Offset);
			}
		}
	}

	SymCleanup(process);
	_writeString(L"\n");
	::operator delete(symbol);
	return ;
}

void CDebugManager::_onAllocation(size_t uSize, 
								  const wchar_t *pwszSymbol,
								  const wchar_t *pwszFileName,
								  UINT uLineNumber,
								  void *pAddress)
{
	if(pwszSymbol == NULL || pwszFileName == NULL) return ;
	AllocListElement *p = cnew AllocListElement;
	if(p == NULL) return ;

	wcsncpy(p->symbol, pwszSymbol, 32);
	p->line = uLineNumber;
	p->size = uSize;

	const wchar_t *pStr = pwszFileName + wcslen(pwszFileName) - 1;
	while(pStr != pwszFileName) {
		if(*pStr == '\\' || *pStr == '/') break;
		pStr--;
	}
	wcsncpy(p->source, pStr, 32);
	
	_pAllocList[pAddress] = p;
}

void CDebugManager::_onDeallocation(void *pAddress)
{
	auto iter = _pAllocList.find(pAddress);
	if(iter != _pAllocList.end()) {
		delete iter->second;
		_pAllocList.erase(iter);
	}
}

void CDebugManager::printMemoryStatus(void)
{
	size_t allocated = 0, count = 0;
	_writeString(L"\n Symbol Name      Source File Name                 Line  Size\n");
	if(_pAllocList.empty()) {
		_writeString(L"  [empty]\n\n");
	} else {
		std::for_each(_pAllocList.begin(), _pAllocList.end(), [&](std::pair<void *, AllocListElement *> elem)
			{
				printRawString(L" %-16.16s %-32.32s %-5d %d\n",
					elem.second->symbol, elem.second->source,
					elem.second->line, elem.second->size);
				allocated += elem.second->size, count++;
			});
		_writeString(L"\n");
	}
	printRawString(L" Total %d bytes unfreed, %d objects.\n\n", allocated, count);
}

void CDebugManager::_freeAllocList(void)
{
	auto iter = _pAllocList.begin();
	for( ; iter != _pAllocList.end(); ++iter) {
		delete iter->second;
	}
	_pAllocList.clear();
}

